"""
A general implementation of genetic algorithm (GA).

GA works by applying three evolutionary operators, mutation, selection, and
crossover, to a randomly-generated "population", or list of "chromosomes"
(potential solutions represented by a NumPy array of integers or floats),
until the minimum solution "fitness" has been reached, or the maximum
number of generations has been exceeded.
In addition, an optional adaptation operator is defined that can be used to
dynamically adapt the magnitude of mutations and the population size to
improve the efficacy and performance of the algorithm.

Most functions in this module are higher-order functions that return another
function. To create a GA optimizer, you must minimally call create_optimizer
with fitness, mutate, and crossover functions. The mutate and crossover
functions are typically created with functions defined by this module
(e.g. create_generic_mutate, create_point_crossover), though you may optionally
define your own. You must supply the fitness function yourself.

Copyright (C) 2013 Jerrad Michael Genson

This program is free software: you can redistribute it and/or modify
it under the terms of the BSD 3-Clause License as published by
the Open Source Initiative.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
BSD 3-Clause License for more details.

You should have received a copy of the BSD 3-Clause License
along with this program.  If not, see
<https://betelgeuse-ml.googlecode.com/hg/LICENSE>

"""

# Python built-in modules
import random
import math
import copy

# Third party modules
import numpy as np

# Betelgeuse modules
import optimize

# Module header
__author__ = 'Jerrad Genson'
__contact__ = 'jerradgenson@gmail.com'
__copyright__ = 'Copyright 2013 Jerrad Michael Genson'
__license__ = 'BSD 3-Clause'

# Global constants
# The maximum mutation magnitude that adaptation might choose.
MAX_MUTATION_RATE = 0.9

# The minimum mutation magnitude that adaptation might choose.
MIN_MUTATION_RATE = 0.1

# The maximum percentage that adaptation can reduce the population by.
MAX_POPULATION_DECREASE = 0.2

# The minimum population size that adaptation can choose.
MIN_POPULATION_SIZE = 25

# Specifies how often to run the adaptor per cycle of GA.
ADAPTATION_RATE = 10


def group_chromosomes(chromosomes, fitness):
    """
    Organize chromosomes into groups as represented by the ChromosomeGroup class.
    A ChromosomeGroup contains the number of chromosomes in that group (i.e. the
    number of duplicates), the fitness for that chromosome, and the
    chromosome itself.

    Args:
        chromosomes: A list of Numpy Arrays representing chromosomes.
        fitness: A fitness function that takes a chromosome and returns
                 its fitness as a real value between 0 and 1.

    Returns:
        A list of ChromosomeGroup instances.

    """

    groups = {}
    # Put every chromosome into a group.
    for chromosome in chromosomes:
        hash_value = optimize.array_hash(chromosome)
        try:
            groups[hash_value].increment()
            # Chromosome matches this group.

        except KeyError:
            # Create a new group for this chromosome.
            groups[hash_value] = ChromosomeGroup(chromosome, fitness(chromosome))

    return list(groups.values())


def create_tournament_select():
    """
    Create a selection operator function using the tournament selection algorithm.
    Tournament selection works by choosing a few chromosomes from the population
    at random and having them "compete" in a tournament. The chromosome with the
    greatest fitness wins. This process is repeated until the target keep value
    has been reached. Chromosomes are discarded from the initial population upon
    selection.

    Returns:
        A selection function.

    """

    def tournament_select(chromosomes, fitness, keep):
        """
        GA selection operator implemented using tournament selection.
        This implementation of tournament selection discards chromosomes from
        the original population after selecting them for preservation.

        Args:
            chromosomes: A list of chromosomes (Numpy Arrays)
            fitness: A function that accepts a chromosomes as input and
                     outputs its fitness.
            keep: The number of chromosomes to keep.

        Returns:
            A list of chromosomes that survived selection.

        """

        # Remember selected chromosomes.
        selection = []
        # Calculate number of chromosomes in each tournament.
        tournament_size = int(math.ceil(len(chromosomes) / float(keep)))
        # Group chromosomes together.
        groups = group_chromosomes(chromosomes, fitness)
        # Run one tournament for every chromosomes we want to keep.
        for tournament_number in range(keep):
            # Choose chromosomes at random to compete in tournament.
            tournament = []
            # Choose chromosomes without repeats.
            while len(tournament) < tournament_size and groups:
                group = random.choice(groups)
                tournament.append(group)
                group.decrement()
                if group.count == 0:
                    # No more chromosomes in this group.
                    groups.remove(group)

            # Select fittest chromosomes in tournament for survival.
            survivor = max(tournament, key=lambda group: group.fitness)
            selection.append(survivor.chromosome)

        return selection

    return tournament_select


def create_roulette_select():
    """
    Create a selection function using the roulette-wheel selection algorithm.
    In roulette-wheel selection, each chromosome has an equal chance of
    being considered for selection. When a chromosome is considered, a
    random real number between 0 and 1 is generated, and the chromosome is
    selected if its fitness is higher than that number. This process is
    repeated until the target keep value has been reached. Once a chromosome
    is selected, it is removed from the initial population.

    Returns:
        A selection function.

    """

    def roulette_select(chromosomes, fitness, keep):
        """
        GA selection operator implemented using roulette-wheel selection
        (also known as fitness proportionate selection).

        Args:
            chromosomes: A list of chromosomes (Numpy Arrays)
            fitness: A function that accepts a chromosomes as input and
                     outputs its fitness.
            keep: The number of chromosomes to keep.


        Returns:
            A list of chromosomes that survived selection.

        """

        # Group chromosomes together.
        groups = group_chromosomes(chromosomes, fitness)
        survivors = []
        while len(survivors) < keep:
            # Ensure each chromosome has an equal chance of being considered.
            random.shuffle(groups)
            # Look at each chromosome and use its fitness as the probability
            # that it will be preserved.
            for group in groups:
                if group.fitness >= random.random():
                    survivors.append(group.chromosome)
                    if len(survivors) == keep:
                        break

        return survivors

    return roulette_select


def create_truncate_select():
    """
    Create a selection function using the truncation selection algorithm.
    In truncation selection, the fittest chromosomes are straightforwardly
    selected for survival. The keep value determines the number of chromosomes
    that are selected.

    Returns:
        A selection function.

    """

    def truncate_select(chromosomes, fitness, keep):
        """
        GA selection operator implemented using truncate selection.

        Args:
            chromosomes: A list of chromosomes (Numpy Arrays)
            fitness: A function that accepts a chromosomes as input and
                     outputs its fitness.
            keep: The number of chromosomes to keep.

        Returns:
            A list of chromosomes that survived selection.

        """

        selection = []
        while chromosomes and len(selection) < keep:
            survivor = optimize.find_best(chromosomes, fitness)[0]
            selection.append(survivor)
            chromosomes = optimize.remove_array(survivor, chromosomes)

        return selection

    return truncate_select


def create_adaptor(fitness,
                   max_mutation_rate=MAX_MUTATION_RATE,
                   min_mutation_rate=MIN_MUTATION_RATE,
                   max_population_decrease=MAX_POPULATION_DECREASE,
                   min_population_size=MIN_POPULATION_SIZE,
                   mutation_adaptation=True,
                   population_adaptation=True,
                   adaptation_rate=ADAPTATION_RATE):
    """
    Create an adaptor function for use with genetic algorithm.
    The adaptation algorithm is a novel solution to the problem of choosing
    good population sizes and mutation magnitudes. It varies both values
    dynamically as GA runs, making adjustments to maximize solution quality.
    The idea is to start with an "exploration phase", with a large mutation
    rate and population size, and decrease both as the fitness increases to
    transition to an "exploitation phase" that converges quickly.

    Args:
        fitness: A function that accepts a chromosome as input and returns
                 its fitness as a real number between 0 and 1.

    Keywords:
        max_mutation_rate: The maximum mutation magnitude that adaptation
                           might choose.
        min_mutation_rate: The minimum mutation magnitude that adaptation
                           might choose.
        max_population_decrease: The maximum percentage that adaptation can
                                 reduce the population by.
        min_population_size: The minimum population size that adaptation can
                             choose.
        mutation_adaptation: True to adapt mutation size, or False to not.
        population_adaptation: True to adapt population size, or False to not.
        adaptation_rate: Specifies how many cycles of GA to complete
                            before running the adaptation algorithms.

    Returns:
        A synthesized adapt function - see documentation for adapt.

    """

    # Keeps track of number of times the adaptation functions has been called.
    times_called = -1

    def adapt_mutation(median_fitness):
        """ Adapt genetic algorithm's mutation magnitude. """

        # Calculate ratio to reduce the mutation rate by inverse proportionally
        # to the median fitness.
        reduction_ratio = 1.0 - median_fitness
        # Reduce the mutation rate by the reduction ratio.
        new_mutation_rate = max_mutation_rate * reduction_ratio
        # Ensure the new mutation rate isn't less than the lower bound.
        if new_mutation_rate < min_mutation_rate:
            return min_mutation_rate

        else:
            return new_mutation_rate

    def adapt_population(median_fitness, old_population_size):
        """ Adapt genetic algorithm's population size. """

        # Calculate ratio to reduce the population size by inverse proportionally
        # to the median fitness. But do not use a higher reduction ratio than
        # specified by max_population_decrease.
        reduction_ratio = 1.0 - median_fitness * max_population_decrease
        # Reduce the population size by the reduction ratio.
        new_population_size = old_population_size * reduction_ratio
        # Ensure the new population size isn't less than the lower bound.
        if new_population_size < min_population_size:
            return min_population_size

        else:
            return new_population_size

    def adapt(chromosomes, mutation_rate):
        """
        Standard genetic algorithm adaptation function.

        Args:
            chromosomes: The current population of chromosomes.
            mutation_rate: The size of the current mutations.

        Returns:
            A tuple of (new mutation size, new population size).

        """

        nonlocal times_called

        times_called += 1
        if times_called % adaptation_rate == 0:
            # Run adaptation this cycle.
            # Find median fitness of chromosomes.
            fitness_list = sorted((fitness(chromosome) for chromosome in chromosomes))
            middle_index = int(round(len(fitness_list) / 2))
            median_fitness = fitness_list[middle_index]

            if mutation_adaptation:
                new_mutation_size = adapt_mutation(median_fitness)

            else:
                new_mutation_size = mutation_rate

            if population_adaptation:
                new_population_size = adapt_population(median_fitness, len(chromosomes))

            else:
                new_population_size = len(chromosomes)

            return new_mutation_size, new_population_size

        else:
            # Do not run adaptation this cycle.
            return mutation_rate, len(chromosomes)

    return adapt


def create_optimizer(fitness,
                     mutate,
                     crossover,
                     select=create_tournament_select(),
                     adapt=None,
                     mutation_rate=MIN_MUTATION_RATE,
                     elitism=False,
                     *args):
    """
    Create a genetic algorithm optimizer function.

    Args:
        fitness: GA fitness operator.
                 A function that takes a chromosome and returns its
                 fitness as a real number between 0 and 1.
        mutate: GA mutation operator.
                A function that takes a chromosome and a mutation
                magnitude (real number between 0 and 1) and mutates the
                chromosome by that magnitude, returning the mutated
                chromosome.
        crossover: GA crossover operator.
                   A function that takes two parent chromosomes and
                   returns a hybrid child chromosome.
        select: GA selection operator.
                A function that takes as inputs a chromosome list,
                a fitness function, and the number of chromosomes to keep
                and returns a reduced ("select") list of chromosomes.

    Keywords:
        adapt: GA adaptation operator.
               A function that take a chromosome list as input and
               returns tuple of (new mutation size, new population size).
        mutation_rate: A static mutation magnitude to use if mutation
                       adaptation is not used. Must be a real number
                       between 0 and 1.
        elitism: Set to True to always preserve the fittest chromosome
                 verbatim through selection, crossover, and mutation.
                 This will have the side-effect of increasing population
                 size by one.

    *args:
        Any number of post-processing functions that take a list of
        chromosomes as input and return a list of chromosomes as output,
        and run at the end of every GA cycle.

    Returns:
        A genetic algorithm optimizer function. See documentation for
        ga_optimizer.

    """

    def ga_optimizer(chromosomes):
        """
        A genetic algorithm optimizer function.

        Args:
            chromosomes A randomly-generated population of chromosomes to
                        optimize.

        Returns:
            The chromosome population after applying one cycle of GA.

        """

        nonlocal mutation_rate

        population_size = len(chromosomes)
        if adapt:
            # Apply adaptation operator.
            mutation_rate, population_size = adapt(chromosomes, mutation_rate)

        # Selection operator should keep a square root of the new population
        # size, because GA will use crossover to square the population size.
        keep = int(math.ceil(math.sqrt(population_size)))

        # Apply selection operator.
        selected_chromosomes = select(chromosomes, fitness, keep)

        # Apply crossover operator.
        # Cross every chromosome with every other chromosome.
        child_chromosomes = []
        for parent1 in selected_chromosomes:
            for parent2 in selected_chromosomes:
                child = crossover(parent1, parent2)
                # Apply mutation operator.
                mutated_child = mutate(child, mutation_rate)
                child_chromosomes.append(mutated_child)

        # Apply post-processing functions.
        for process in args:
            child_chromosomes = process(child_chromosomes)

        if elitism:
            # Identify fittest parent chromosome and append it to
            # child_chromosomes.
            elite = optimize.find_best(chromosomes, fitness)[0]
            child_chromosomes.append(elite)

        return child_chromosomes

    return ga_optimizer


def create_generic_mutate(bounds=None):
    """
    Create a generic mutation function.
    Makes a change to a number of genes (determined by magnitude) in the
    chromosome. The maximum amount of change that may be applied to a
    gene is also determined by the mutation rate.

    Args:
        bounds: A range of values that genes can take on. Specified as a
                tuple of (low bound, high bound).

    Returns:
        A generic_mutate function. See documentation for generic_mutate.

    """

    def generic_mutate(chromosome, mutation_rate):
        """
        A generic genetic algorithm mutation function.
        Makes a change to a number of genes (determined by mutation_rate) in the
        chromosome. The maximum amount of change that may be applied to a
        gene is also determined by mutation_rate.

        Args:
            chromosome: A Numpy Array of integer or real numbers.
                        Copied before any changes are made.
            mutation_rate: The amount of mutation to apply as a real number
                           between 0 and 1, where 0 is no mutation, and 1 is
                           the most mutation possible.

        Returns:
            A chromosome that has been mutated.

        """

        chromosome_copy = np.copy(chromosome)
        # Calculate the number of mutations to make.
        mutations = int(round(len(chromosome) * mutation_rate))
        # Now apply the mutations.
        for mutation in range(mutations):
            # Randomly decide if the change will be positive or negative.
            sign = random.choice((1, -1))
            # Randomly choose a gene to mutate.
            highest_index = len(chromosome) - 1
            gene_index = random.randint(0, highest_index)
            # Calculate the change to apply based on sign and magnitude.
            # Add 1 to ensure change is nonzero.
            change = sign * mutation_rate * chromosome_copy[gene_index] + 1 * sign
            if chromosome.dtype == 'int':
                # Make sure change is an integer.
                change = int(math.ceil(change))

            # Apply the mutation.
            new_gene = chromosome_copy[gene_index] + change
            # Make sure the new gene is within specified boundaries.
            if bounds and new_gene < bounds[0]:
                new_gene = bounds[0]

            elif bounds and new_gene > bounds[1]:
                new_gene = bounds[1]

            chromosome_copy[gene_index] = new_gene

        return chromosome_copy

    return generic_mutate


def create_shuffle_mutate():
    """
    Create a a mutation function using the shuffle mutation algorithm.
    A mutation algorithm that shuffles genes in a chromosome instead of
    modifying them in place.

    Returns:
        A shuffle_mutate function. See documentation for shuffle_mutate.

    """

    def shuffle_mutate(chromosome, mutation_rate):
        """
        A mutation algorithm that shuffles genes in a chromosome instead of
        actually modifying them.

        Args:
            chromosome: A Numpy Array of integer or real numbers.
                        Copied before any changes are made.
            mutation_rate: The amount of mutation to apply as a real number
                           between 0 and 1, where 0 is no mutation, and 1 is
                           the most mutation possible.

        Returns:
            A shuffle-mutated numpy array based on the given chromosome.

        """

        # Calculate the number of mutations to make.
        mutations = int(round(len(chromosome) * mutation_rate / 2))
        # Find the index of the last gene in the chromosome.
        highest_index = len(chromosome) - 1
        # Now apply the mutations.
        for mutation in range(mutations):
            # Randomly choose two indices to swap.
            index1 = random.randint(0, highest_index)
            index2 = random.randint(0, highest_index)
            # Swap the elements at the two chosen indexes.
            chromosome = optimize.swap(index1, index2, chromosome)

        return chromosome

    return shuffle_mutate


def create_null_mutate():
    """
    Create a null mutation function.
    A mutation function that does nothing but return the unaltered
    chromosome that was given as input.

    Returns:
        A null_mutate function. See documentation for null_mutate for details.

    """

    # noinspection PyUnusedLocal
    def null_mutate(chromosome, mutation_rate):
        """
        A mutation function that does nothing but return the unaltered
        chromosome that was given as input.

        Args:
            chromosome Polymorphic - can be any value. Not used.
            mutation_rate Polymorphic - can be any value.

        Returns:
            The unaltered chromosome argument.

        """

        return chromosome

    return null_mutate


def create_point_crossover():
    """
    Create a single-point crossover function.
    A single point crossover function. The genetic algorithm crossover
    operator is used to produce a child chromosome from the combination
    of two parent chromosomes. This implementation of crossover works by
    splicing the first half of the first parent chromosome with the second
    half of the second parent chromosome to form the hybrid child
    chromosome.

    Returns:
        A point_crossover function. See documentation for point_crossover.

    """

    def point_crossover(chromosome1, chromosome2):
        """
        A single point crossover function. The genetic algorithm crossover
        operator is used to produce a child chromosome from the combination
        of two parent chromosomes. This implementation of crossover works by
        splicing the first half of the first parent chromosome with the
        second half of the second parent chromosome to form the hybrid child
        chromosome.

        Args:
            chromosome1: The first parent chromosome. Must support array slicing.
            chromosome2: The second parent chromosome. Must support array slicing.

        Returns:
            A child chromosome represented by a Numpy Array.

        """

        # Isolate first half of chromosome1.
        split_index1 = int(math.ceil(len(chromosome1) / 2))
        first_half = chromosome1[:split_index1]
        # Isolate second half of chromosome2.
        split_index2 = int(math.ceil(len(chromosome2) / 2))
        second_half = chromosome2[split_index2:]
        # Create child chromosome by concatenating first half of first parent
        # and second half of second parent.
        child = np.concatenate((first_half, second_half))

        return child

    return point_crossover


def create_random_crossover():
    """
    Create a random crossover function.
    A crossover algorithm that takes genes from parent chromosomes at random
    instead of in a predefined sequence. Genes in the child will retain the
    order of the genes in the parents, however, each gene will be chosen at
    random from one of the two corresponding parent genes. Parent
    chromosomes need not be the same length to use this function (if the
    lengths differ, the algorithm will decide at random whether to complete
    the child with genes from the longer chromosome or not).

    Returns:
        A random_crossover function. See documentation for random_crossover.

    """

    def random_crossover(chromosome1, chromosome2):
        """
        A crossover algorithm that takes genes from parent chromosomes at
        random instead of in a predefined sequence. Genes in the child will
        retain the order of the genes in the parents, however, each gene
        will be chosen at random from one of the two corresponding parent
        genes. Parent chromosomes need not be the same length to use this
        function (if the lengths differ, the algorithm will decide at
        random whether to complete the child with genes from the longer
        chromosome or not).

        Args:
            chromosome1: The first parent chromosome, represented by a numpy
                         array.
            chromosome2: The second parent chromosome, represented by a numpy
                         array.

        Returns:
            A child chromosome that is a hybrid of chromosome1 and chromosome2.

        """

        # Choose a chromosome to iterator over at random so if they are not the
        # same length, the child's length will be determined at random.
        if random.choice((True, False)):
            chromosome_iterator = chromosome1
            other_chromosome = chromosome2

        else:
            chromosome_iterator = chromosome2
            other_chromosome = chromosome1

        # Create an empty child chromosome.
        child = np.empty([len(chromosome_iterator)],
                         dtype=chromosome_iterator.dtype)

        for index, gene in enumerate(chromosome_iterator):
            if random.choice((True, False)):
                # Use gene from chromosome_iterator.
                child[index] = gene

            else:
                # Use gene from other_chromosome.
                try:
                    child[index] = other_chromosome[index]

                except IndexError:
                    # chromosome_iterator is longer than other_chromosome.
                    # Use gene from chromosome_iterator instead.
                    child[index] = gene

        return child

    return random_crossover


def create_null_crossover():
    """
    Create a null crossover function.
    A crossover function that does nothing but return an unaltered chromosome.

    Returns:
        A null_crossover function. See documentation for null_crossover.

    """

    # noinspection PyUnusedLocal
    def null_crossover(chromosome1, chromosome2):
        """
        A crossover function that does nothing but return an unaltered chromosome.

        Args:
            chromosome1: Polymorphic - can be any value.
            chromosome2: Polymorphic - can be any value.

        Returns:
            The unaltered chromosome1 argument.

        """

        return chromosome1

    return null_crossover


class ChromosomeGroup:
    """
    Used to group duplicate chromosomes. Has attributes for the chromosome
    value itself, the fitness of the chromosome, and the number of times it
    is duplicated.

    Args:
        chromosome The actual chromosome (usually a Numpy array).
        fitness: The fitness of the chromosome.

    """

    def __init__(self, chromosome, fitness):
        self.__chromosome = chromosome
        self.__fitness = fitness
        self.__count = 1

    @property
    def chromosome(self):
        """ The actual chromosome (usually a Numpy array). """

        return copy.copy(self.__chromosome)

    @property
    def fitness(self):
        """ The fitness of the chromosome. """

        return self.__fitness

    @property
    def count(self):
        """ The number of times the chromosome is duplicated. """

        return self.__count

    def increment(self):
        """ Increment the ChromosomeGroup count by 1. """

        self.__count += 1

    def decrement(self):
        """ Decrement the ChromosomeGroup count by 1. """

        self.__count -= 1
